# -*- coding: utf-8 -*-
"""
Created on Thu Aug 29 14:15:36 2024

@author: Administrator
"""

import tensorflow as tf


def check_tensorflow_gpu():
    print("TensorFlow 版本:", tf.__version__)
    # 检查tensorflow是否可以获取到GPU
    if tf.config.list_physical_devices('GPU'):
        print("GPU is available")
        gpus = tf.config.list_physical_devices('GPU') # 检查GPU设备列表
        for gpu in gpus:
            print(gpu, "设备类型:", gpu.device_type, "设备名称:", gpu.name)
            tf.config.experimental.set_memory_growth(gpu, True)    # 开启内存增长。此选项会尝试根据运行时分配的需求分配尽可能充足的GPU内存
        tf.config.set_visible_devices(gpus[0], 'GPU') # 使用第一个GPU
    else:
        print("GPU is not available, using CPU")
    # 检查tensorflow是否得到CUDA支持
    if tf.test.is_built_with_cuda():
        print("CUDA is available")


if __name__ == "__main__":
    check_tensorflow_gpu()
    


"""
PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU') 设备类型: CPU 设备名称: /physical_device:CPU:0
PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU') 设备类型: GPU 设备名称: /physical_device:GPU:0
"""